import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def get_train_data(data_size=100):
    data_label = np.zeros((2 * data_size, 1))
    # class 1
    x1 = np.reshape(np.random.normal(1, 0.6, data_size), (data_size, 1))
    y1 = np.reshape(np.random.normal(1, 0.8, data_size), (data_size, 1))
    data_train = np.concatenate((x1, y1), axis=1)
    data_label[0:data_size, :] = 0
    # class 2
    x2 = np.reshape(np.random.normal(-1, 0.3, data_size), (data_size, 1))  # 0.3
    y2 = np.reshape(np.random.normal(-1, 0.5, data_size), (data_size, 1))  # 0.5
    data_train = np.concatenate((data_train, np.concatenate((x2, y2), axis=1)), axis=0)
    data_label[data_size: 2 * data_size, :] = 1
    return data_train, data_label


def get_test_data(data_size=10):
    testdata_label = np.zeros((2 * data_size, 1))
    # class 1
    x1 = np.reshape(np.random.normal(1, 0.6, data_size), (data_size, 1))
    y1 = np.reshape(np.random.normal(1, 0.8, data_size), (data_size, 1))
    data_test = np.concatenate((x1, y1), axis=1)
    testdata_label[0:data_size, :] = 0
    # class 2
    x2 = np.reshape(np.random.normal(-1, 0.3, data_size), (data_size, 1))  # 0.3
    y2 = np.reshape(np.random.normal(-1, 0.5, data_size), (data_size, 1))  # 0.5
    data_test = np.concatenate((data_test, np.concatenate((x2, y2), axis=1)), axis=0)
    testdata_label[data_size: 2 * data_size, :] = 1
    return data_test, testdata_label


def linear_regression(train_dat, train_label, test_dat, test_label):
    train_size, test_size = int(train_label.size / 2), int(test_label.size / 2)
    ax = plt.axes(projection='3d')

    # change label from (1/0) to (+1/-1)
    np.place(train_label, train_label == 0, -1)
    np.place(test_label, test_label == 0, -1)

    # regularization
    train_dat = np.c_[train_dat, np.ones((train_dat.shape[0], 1))]

    # calculate weight
    weight = np.linalg.inv(train_dat.T.dot(train_dat)).dot(train_dat.T).dot(train_label)

    # draw sample points
    ax.scatter3D(train_dat[:train_size, 0], train_dat[:train_size, 1], train_label[:train_size],
                 label='class 1 train data')
    ax.scatter3D(train_dat[train_size:, 0], train_dat[train_size:, 1], train_label[train_size:],
                 label='class 2 train data')

    # draw regression surface
    X, Y = np.meshgrid(np.linspace(-2, 2, 100), np.linspace(-2, 2, 100))
    ax.plot_surface(X, Y, weight[0] * X + weight[1] * Y + weight[2])

    # draw test data
    ax.scatter3D(test_dat[:test_size, 0], test_dat[:test_size, 1], test_label[:test_size], c='r', marker='v',
                 label='class 1 test data')
    ax.scatter3D(test_dat[test_size:, 0], test_dat[test_size:, 1], test_label[test_size:], c='k', marker='v',
                 label='class 2 test data')
    plt.legend(loc=3)
    plt.title('Result of Linear Regression: same variance')
    plt.savefig('lr_same.pdf')


def linear_discriminant_analysis(train_dat, train_label, test_dat, test_label):
    train_size, test_size = int(train_label.size / 2), int(test_label.size / 2)

    # calculate statistics
    data_1, data_2 = train_dat[:train_size, :].T, train_dat[train_size:, :].T
    test_1, test_2 = test_dat[:test_size, :].T, test_dat[test_size:, :].T
    mu_1, mu_2 = data_1.mean(axis=1).reshape(2, 1), data_2.mean(axis=1).reshape(2, 1)
    sigma_1, sigma_2 = (data_1 - mu_1).dot((data_1 - mu_1).T), (data_2 - mu_2).dot((data_2 - mu_2).T)

    # calculate weight
    weight = np.linalg.inv(sigma_1 + sigma_2).dot(mu_1 - mu_2)
    print(weight)

    # draw train data
    plt.scatter(data_1[0, :], data_1[1, :], label='class 1 train data')
    plt.scatter(data_2[0, :], data_2[1, :], label='class 2 train data')

    # draw the projection of train data
    projection_1 = weight.dot(weight.T.dot(data_1)) / np.dot(weight.T, weight)
    projection_2 = weight.dot(weight.T.dot(data_2)) / np.dot(weight.T, weight)
    plt.scatter(projection_1[0], projection_1[1], label='class 1 train projection')
    plt.scatter(projection_2[0], projection_2[1], label='class 2 train projection')

    # draw test data
    plt.scatter(test_1[0, :], test_1[1, :], label='class 1 test data', marker='^')
    plt.scatter(test_2[0, :], test_2[1, :], label='class 2 test data', marker='^')

    # &&&&&&&&&&& TEST DATA CLASSIFICATION &&&&&&&&&&&
    threshold = (weight.T.dot(mu_1) + weight.T.dot(mu_2)) / 2
    projection_test_1 = weight.T.dot(test_1)
    projection_test_2 = weight.T.dot(test_2)
    ACCURACY = (np.sum(projection_test_1 > threshold) + np.sum(projection_test_2 < threshold)) / (test_size * 2)
    print(ACCURACY)

    # draw the classification line
    threshold_y = -threshold * weight[0] / weight[1]
    X = np.linspace(-1, 1, 100)
    plt.plot(X, (-weight[0] / weight[1] * (X - threshold) + threshold_y).flatten(), label='classification line')

    # draw the projection of test data
    projection_test_1 = weight.dot(projection_test_1) / np.dot(weight.T, weight)
    projection_test_2 = weight.dot(projection_test_2) / np.dot(weight.T, weight)
    plt.scatter(projection_test_1[0], projection_test_1[1], label='class 1 test projection', marker='v', c='r')
    plt.scatter(projection_test_2[0], projection_test_2[1], label='class 2 test projection', marker='v', c='k')

    plt.legend()
    plt.title('Result of LDA: different variance; ACCURACY={}'.format(ACCURACY))
    plt.axis('equal')
    plt.savefig('lda_diff.pdf')


if __name__ == '__main__':
    train_data, train_data_label = get_train_data()
    test_data, test_data_label = get_test_data()
    # linear_regression(train_data, train_data_label, test_data, test_data_label)
    linear_discriminant_analysis(train_data, train_data_label, test_data, test_data_label)
